# coding=utf-8
# @Time    : 2021/2/3 14:09
# @Author  : Leo
# @Email   : l1512102448@qq.com
# @File    : image_util.py


import numpy as np
import PIL.Image as Image
import PIL
# import matplotlib.pyplot as plt

def pil_to_np(img_PIL):
    '''Converts image in PIL format to np.array.
    From W x H x C [0...255] to C x W x H [0..1]
    '''
    ar = np.array(img_PIL)

    if len(ar.shape) == 3:
        pass
        # ar = ar.transpose(2, 0, 1) # WxHx3 --> 3xWxH
    else:
        ar = ar[None, ...]  # WxH --> 1xWxH

    return ar.astype(np.float32)/255.

def np_to_pil(img_np):
    '''Converts image in np.array format to PIL image.
    From C x W x H  to  W x H x C [0...255]
    '''
    if np.max(img_np[:]) <= 1:
        ar = np.clip(img_np * 255, 0, 255).astype(np.uint8)
    else:
        ar = np.clip(img_np, 0, 255).astype(np.uint8)
    if img_np.shape[0] == 1:  # 灰度图 1xWxH --> WxH
        ar = ar[0]
    else:  # 彩色图 要转成   3xWxH -->  WxHx3
        ar = ar.transpose(1, 2, 0)

    return Image.fromarray(ar)

def get_image(path, imsize=-1):
    """Load an image and resize to a cpecific size.
    读取图像
    Args:
        path: path to image
        imsize: tuple or scalar with dimensions; -1 for `no resize`
    """
    img = Image.open(path)
    img = img.convert("RGB")
    if isinstance(imsize, int):
        imsize = (imsize, imsize)

    if imsize[0]!= -1 and img.size != imsize:  # resize
        if imsize[0] > img.size[0]:
            img = img.resize(imsize, Image.BICUBIC)
        else:
            img = img.resize(imsize, Image.ANTIALIAS)

    img_np = pil_to_np(img)

    return img, img_np

# 图像旋转变换 8种
def data_augmentation(I, K):

    '''
    :param I: PIL image
    :param K: 1,2,3,4,..8
    :return: 对图像进行 旋转或翻转 或 旋转+翻转
    '''
    if K == 1:
        return I
    elif K == 2:
        return I.transpose(PIL.Image.FLIP_TOP_BOTTOM)  #上下翻转
    elif K == 3:
        I = I.rotate(90)  # 旋转
        return I
    elif K == 4:
        I = I.rotate(90)
        return I.transpose(PIL.Image.FLIP_TOP_BOTTOM)
    elif K == 5:
        I = I.rotate(180)
        return I
    elif K == 6:
        I = I.rotate(180)
        return I.transpose(PIL.Image.FLIP_TOP_BOTTOM)
    elif K == 7:
        I = I.rotate(270)
        return I
    elif K == 8:
        I = I.rotate(270)
        return I.transpose(PIL.Image.FLIP_TOP_BOTTOM)


